import mediapipe as mp
import cv2
import os
import numpy as np
import csv


# 这个里面写的就是根据label来获得对应的landmark
# 如果后续用的话这个函数要改下，改为传参版本，不用改里面的
def google_landmark_label(origin_label, image_folder, csv_folder):
    mp_face = mp.solutions.face_mesh
    face_mesh = mp_face.FaceMesh(
          static_image_mode=True,   # 是否在静态图像中运行
          max_num_faces=1,
          refine_landmarks=True,    # 优化关键点的位置
          min_detection_confidence=0.5,  # 人脸置信度
          min_tracking_confidence=0.5   # 跟踪人脸的最小置信度
    )
    with open(origin_label, 'r') as label_file:
        lines = label_file.readlines()
        # lines.pop(0)
    # iii=1
    for line in lines:
        # print(iii)
        # iii+=1
        # 对于mpii是p00\1.jpg p00\left\1.jpg p00\right\1.jpg day08/0069.jpg left -0.252647142956226,0.05464552625489694,-0.966014123921789 -0.46732565061917913,-0.05736579094793633,-0.013657707986779405 0.2558059438789034,-0.05467275933864655 -0.058433881176183707,0.4674596468469367 0.01715928,-0.10456196,-0.037469573 1.0,1.0,1.240293724425897 2.1582156080057757e-06,9.124974553742504e-08,599.9999396334617
        image_name = line.strip().split(" ")[0]  # 比如说这里获得的就是p00\1.jpg
        folder_name = image_name.split("\\")[0]     # p00
        image_jpg = image_name.split("\\")[1]          # 1.jpg
        image_path = os.path.join(image_folder, folder_name, image_jpg)  # 原图路径
        image = cv2.imread(image_path)
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 将图像传递给MediaPipe Face模型进行处理
        results = face_mesh.process(image_rgb)
        eye_landmarks = []
        # 获取眼睛关键点坐标
        if results.multi_face_landmarks:
            for face_landmarks in results.multi_face_landmarks:
                # 等于说face_landmarks是一个人的脸部关键点

                # 这下面是添加.label信息
                for i in range(7):         # mpiigaze是12个   360是6
                    eye_landmarks.append(line.strip().split(" ")[i])
                for i in range(478):
                    eye_landmarks.append(np.asarray(face_landmarks.landmark[i].x * image.shape[0]))
                    eye_landmarks.append(np.asarray(face_landmarks.landmark[i].y * image.shape[0]))

            # 将关键点追加到csv文件
            with open(csv_folder, mode='a', newline='') as csvfile:
                csv_writer = csv.writer(csvfile)
                # 写入每一行数据
                # data_row = []
                # for item in eye_landmarks:
                    # data_row.extend(item * image.shape[0])
                csv_writer.writerow(eye_landmarks)
    # Close MediaPipe FaceMesh module
    face_mesh.close()


def google_landmark_label_eyediap(origin_label, image_folder, csv_folder, num_eyediap):
    mp_face = mp.solutions.face_mesh
    face_mesh = mp_face.FaceMesh(
          static_image_mode=True,   # 是否在静态图像中运行
          max_num_faces=1,
          refine_landmarks=True,    # 优化关键点的位置
          min_detection_confidence=0.5,  # 人脸置信度
          min_tracking_confidence=0.5   # 跟踪人脸的最小置信度
    )
    with open(origin_label, 'r') as label_file:
        lines = label_file.readlines() # 读取到了所有的label
    iii = 1
    count = 1
    for line in lines:
        # print(iii)
        iii+=1
        if iii % 100 == 0:
            print(iii)
        image_name = line.strip().split(" ")[0]  # 比如说这里获得的就是p1/face/1.jpg
        image_path = os.path.join(image_folder, image_name)  # 原图路径
        image = cv2.imread(image_path)
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 将图像传递给MediaPipe Face模型进行处理
        results = face_mesh.process(image_rgb)
        eye_landmarks = []
        # 获取眼睛关键点坐标
        if results.multi_face_landmarks:
            count += 1
            print(count)
            for face_landmarks in results.multi_face_landmarks:
                # 等于说face_landmarks是一个人的脸部关键点

                # 这下面是添加.label信息
                for i in range(num_eyediap):         # mpiigaze是12个   360是6
                    eye_landmarks.append(line.strip().split(" ")[i])
                for i in range(478):
                    eye_landmarks.append(np.asarray(face_landmarks.landmark[i].x * image.shape[0]))
                    eye_landmarks.append(np.asarray(face_landmarks.landmark[i].y * image.shape[0]))

            # 将关键点追加到csv文件
            with open(csv_folder, mode='a', newline='') as csvfile:
                csv_writer = csv.writer(csvfile)
                csv_writer.writerow(eye_landmarks)
    # Close MediaPipe FaceMesh module
    face_mesh.close()

if __name__ == "__main__":
    print("111")
    # 这个是mpii的
    # origin_label = r'/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/MPIIGaze/allLabel/all_nohead.label'
    # image_folder = r'/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/MPIIGaze/Image'
    # csv_folder = r'/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/MPIIGaze/google_all_result_csv/google_allImage.csv'
    # google_landmark_label(origin_label, image_folder, csv_folder)

    # 这个是360的
    # origin_label = r'/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/Gaze360/resultLabel/selectLabel.label'
    # image_folder = r'/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/Gaze360/Image'
    # csv_folder = r'/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/Gaze360/google_all_result_csv_360/google_allImage_360.csv'
    # google_landmark_label(origin_label, image_folder, csv_folder)

    # # 这个是eth的
    # image_folder = '/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/alleth/train'
    # csv_folder = '/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/alleth/google_csv'
    # origin_label_path = '/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/alleth/singleLabel'
    # origin_labels = os.listdir(origin_label_path)
    # for origin_label in origin_labels:
    #     name = origin_label.split(".")[0]
    #     label_path = os.path.join(origin_label_path, origin_label)
    #     csv_folder_path = os.path.join(csv_folder, f'{name}.csv')
    #     google_landmark_label(label_path, image_folder, csv_folder_path)


    # 这个是eyediap的
    origin_label = '/home/xian/eyediap_25_2.17/all_noselect_diap.label'
    image_folder = '/home/xian/eyediap_25_2.17/Image'
    csv_folder = '/home/xian/eyediap_25_2.17/all_noselect_diap_2025.csv'
    google_landmark_label_eyediap(origin_label, image_folder, csv_folder, 11)
    # image_folder = r'D:\BaiduNetdiskDownload\FaceBased\eyeDiap\yasuo\Image'
    # csv_folder = '/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/alleth/google_csv'
    # origin_label_path = '/home/xian/mazioshuo/chongxinshixian/5.6-7.2-bak/img/alleth/singleLabel'
    # origin_labels = os.listdir(origin_label_path)
    # for origin_label in origin_labels:
    #     name = origin_label.split(".")[0]
    #     label_path = os.path.join(origin_label_path, origin_label)
    #     csv_folder_path = os.path.join(csv_folder, f'{name}.csv')
    #     google_landmark_label(label_path, image_folder, csv_folder_path)